import datetime
import os
import numpy as np
import cv2 as cv
from PyCmpltrtok.common import *
from PyCmpltrtok.common_np import *
from python_ai.category.TensorRT.tutorial_helper import load_engine, infer
from postprocess import postprocess


if '__main__' == __name__:
    import warnings

    warnings.simplefilter("error")
    warnings.simplefilter("ignore", DeprecationWarning)

    def _main():
        # trt_path = '_models/fcn-resnet101.trt'
        # PRECISION = np.float32

        trt_path = '_models/fcn-resnet101.fp16.engine'
        PRECISION = np.float16

        sep('Load image')
        img_path = '/var/asuspei/my_src/TensorRT/TensorRT.git/quickstart/SemanticSegmentation/input.ppm'
        # img_path = '/var/asuspei/large_data/pic/dog_bird.jpg'
        # img_path = '/var/asuspei/large_data/pic/dog.jpg'
        # img_path = '/var/asuspei/large_data/pic/blurred/swim01.jpg'
        # img_path = '/var/asuspei/large_data/pic/cuty/leg01.jpg'  # too larg
        # img_path = '/var/asuspei/large_data/pic/cuty/leg01.resized.png'  # too large
        # img_path = '/var/asuspei/large_data/pic/cuty/leg01.resized.land.png'
        M = 1
        img = cv.imread(img_path, cv.IMREAD_COLOR)
        # cv.imshow('img', img)
        img_ = img.copy()
        img = img[:, :, ::-1]

        img = uint8_to_flt_by_lut(img)

        # norm start
        check_np_detailed(img, 'img')
        check_np_detailed(img[:, :, 0], 'img[:, :, 0]')
        check_np_detailed(img[:, :, 1], 'img[:, :, 1]')
        check_np_detailed(img[:, :, 2], 'img[:, :, 2]')
        mean = np.array([0.485, 0.456, 0.406], dtype=np.float32)
        std = np.array([0.229, 0.224, 0.225], dtype=np.float32)
        img -= mean
        img /= std
        check_np_detailed(img, 'img')
        check_np_detailed(img[:, :, 0], 'img[:, :, 0]')
        check_np_detailed(img[:, :, 1], 'img[:, :, 1]')
        check_np_detailed(img[:, :, 2], 'img[:, :, 2]')
        # norm end

        img = img.transpose(2, 0, 1)
        input_arr = np.expand_dims(img, axis=0)
        check_np(input_arr, 'input_arr')

        sep('Load model by tutorial helper')
        print(trt_path)
        dt1 = datetime.datetime.now()
        engine = load_engine(trt_path)
        dt2 = datetime.datetime.now()
        duration = dt2 - dt1
        print('duration', duration)

        sep('Run it')
        dt1 = datetime.datetime.now()
        output_arr = infer(engine, input_arr, (1, 1, *input_arr.shape[2:]))
        dt2 = datetime.datetime.now()
        duration = dt2 - dt1
        print('duration', duration)
        postprocess(output_arr, img_, 777)

        # cv.waitKey()
        # cv.destroyAllWindows()

    _main()
